from solvers.solverbase import SolverBase
from solvers.crosswind.crosswindbase import CrosswindBase
from solvers.stabilizations import supg
from solvers.stabilizations import crosswind
from dolfin import *

class BackwardEuler(SolverBase,CrosswindBase):
    'Backward-Euler method with crosswind stabilization.'
    def __init__(self,p=1):
        SolverBase.__init__(self,p)
        CrosswindBase.__init__(self)

    def solve(self,problem):
        mesh = problem.domain

        V = FunctionSpace(mesh,'CG',1)
        VV = VectorFunctionSpace(mesh,'CG',1)
        DG = FunctionSpace(mesh,'DG',0)

        phi_ = interpolate(problem.phi_,V)
        phi = TrialFunction(V)
        varphi = TestFunction(V)

        t = problem.t           
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)
        tau = Function(DG)
        tau.vector()[:] = supg.get_tau(v,eps=supg.diffusivity)

        sigma = Function(DG)
        sigma.vector()[:] = crosswind.get_sigma(v,eps=supg.diffusivity)

        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0
        
        pert = conditional(
                           gt(sqrt(inner(v,v)),1E-16),
                           Constant(1.0),
                           Constant(0.) 
                           ) # a way to realize swith on |v| == 0.
                        
        pMatrix = conditional(
                             gt(sqrt(inner(v,v)),1E-16),
                             Identity(2)-outer(v,v)/inner(v,v),
                             Constant(0.)*Identity(2),
                             ) # I - v X v
                
        a = phi*varphi*dx + dt*inner(v,nabla_grad(phi))*varphi*dx +\
            tau*pert*phi*inner(v,nabla_grad(varphi))*dx +\
            dt*tau*pert*inner(dot(nabla_grad(phi),outer(v,v)),nabla_grad(varphi))*dx+\
            dt*sigma*inner(dot(nabla_grad(phi),pMatrix),nabla_grad(varphi))*dx

        l = phi_*varphi*dx + tau*pert*phi_*inner(v,nabla_grad(varphi))*dx

        A = assemble(a)     
        phi = Function(V)
   
        t = 0.
        count = 0

        while t <= T:
            t += float(dt)
            count += 1

            L = assemble(l) 

            solve(A,phi.vector(),L)
            phi_.assign(phi)
            
            self.update(problem,phi_,t,float(dt),count)

        problem.save_data(self.saveDir,str(self.CFL))

